import os
import django

from django.utils import timezone
from celery import shared_task

from celery_client.app import app as celery_client

os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'zmm.settings')
django.setup()
from functools import wraps
from common.models import Tenant
from load_planer.models import Task, TaskStatus

from django.db import transaction, connection
def with_tenant(celery_task):
    @wraps(celery_task)
    def decor(payload, *args, **kwargs):
        tenant_id = payload['tenant_id']
        tenant = Tenant.objects.get(id=tenant_id)
        prev_schemas = connection.schemas
        connection.set_schema(tenant.schema_name)
        with transaction.atomic():
            results = celery_task(payload, *args, **kwargs)
        connection.set_schema(*prev_schemas)

        return results
    return decor
def retrieve_and_lock_task(id):
    return (
        Task
        .objects
        .select_for_update()
        .filter(id=id)
        .first()
    )

@shared_task(name='internal.do_submit')
@with_tenant
def on_task_submitted(payload):
  
    task = retrieve_and_lock_task(payload['task_id'])
    new_payload = {
        'tenant_id': payload['tenant_id'],
        'task_id': task.id,
        'config': task.config,
        'input_dataset': task.input_dataset,
    }
    print(f'提交任务请求体: {new_payload}')
    celery_client.send_task('zmm.do_compute', args=[new_payload])
    task.status = TaskStatus.SUBMITTED.value
    task.submitted_at = timezone.now()
    task.save()

    print('任务已提交')


@shared_task(name='zmm.do_start')
@with_tenant
def on_task_started(payload):
    task = retrieve_and_lock_task(payload['task_id'])
    task.status = TaskStatus.RUNNING.value
    task.save()
    print('任务开始运行')


@shared_task(name='zmm.do_finish')
@with_tenant

def on_task_finished(payload):
    task = retrieve_and_lock_task(payload['task_id'])
    task.result_dataset = payload['result_dataset']
    task.status = TaskStatus.FINISHED.value
    task.finished_at=timezone.now()
    task.save()
    print('任务已结束')


@shared_task(name='zmm.do_fail')
@with_tenant
def on_task_esrror(payload):
    task = retrieve_and_lock_task(payload['task_id'])
    task.status = TaskStatus.ERROR.value
    task.finished_at=timezone.now()
    task.save()
    print('任务错误')
